import matplotlib.pyplot as plt
import numpy as np
import sqlite3
from datetime import datetime

# 定义函数，用于处理数据
def process_data(rows):
    dates = [datetime.strptime(row[0], '%Y-%m-%d') for row in rows]
    total_market_values = [float(row[1]) if row[1] else None for row in rows]
    circulating_market_values = [float(row[2]) if row[2] else None for row in rows]
    listed_companies = [int(row[3]) if row[3] else None for row in rows]
    average_pe_ratios = [float(row[4]) if row[4] else None for row in rows]

    # 计算流通市值占总市值的比例
    ratios = []
    for total, circulating in zip(total_market_values, circulating_market_values):
        if total and circulating:
            ratio = float(circulating) / float(total)
            ratios.append(ratio)
        else:
            ratios.append(None)

    # 过滤掉 None 值
    valid_dates = []
    valid_total_market_values = []
    valid_circulating_market_values = []
    valid_ratios = []
    valid_listed_companies = []
    valid_average_pe_ratios = []
    for j in range(len(dates)):
        if total_market_values[j] and circulating_market_values[j] and ratios[j] and listed_companies[j] and average_pe_ratios[j]:
            valid_dates.append(dates[j])
            valid_total_market_values.append(total_market_values[j])
            valid_circulating_market_values.append(circulating_market_values[j])
            valid_ratios.append(ratios[j])
            valid_listed_companies.append(listed_companies[j])
            valid_average_pe_ratios.append(average_pe_ratios[j])

    return valid_dates, valid_total_market_values, valid_circulating_market_values, valid_ratios, valid_listed_companies, valid_average_pe_ratios

# 定义函数，用于绘制子图
def plot_subplot(ax, dates, values, label, title, xlabel, ylabel):
    if len(dates) == 1:
        ax.scatter(dates, values, label=label)
    else:
        # 绘制曲线
        ax.plot(dates, values, label=label)
        # 绘制数据点
        ax.scatter(dates, values, color='red', zorder=5)  # zorder=5 确保数据点显示在曲线之上
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.legend()
    # 设置 x 轴范围从 2025 年 1 月 1 日起
    from datetime import datetime
    start_date = datetime(2025, 1, 1)
    ax.set_xlim(left=start_date)

# 定义函数，用于绘制数据变化趋势图
def plot_data_trend():
    try:
        # 连接到数据库
        conn = sqlite3.connect('stock_data.db')
        cursor = conn.cursor()

        # 指定支持中文的字体，这里以 SimHei 为例，你可以根据自己系统中的字体进行调整
        plt.rcParams['font.family'] = 'SimHei'
        # 解决负号显示问题
        plt.rcParams['axes.unicode_minus'] = False

        exchanges = ['上海交易所', '深圳交易所']
        plt.figure(figsize=(16, 12))

        all_valid_dates = []
        all_valid_total_market_values = []
        all_valid_circulating_market_values = []
        all_valid_ratios = []
        all_valid_listed_companies = []
        all_valid_average_pe_ratios = []

        for i, exchange in enumerate(exchanges):
            cursor.execute('SELECT date, total_market_value, circulating_market_value, listed_companies, average_pe_ratio FROM stock_data WHERE exchange =?', (exchange,))
            rows = cursor.fetchall()

            if not rows:
                print(f"未找到 {exchange} 的数据")
                continue

            # 打印处理前的数据内容
            print(f"处理前 {exchange} 的数据内容:")
            for row in rows:
                print(row)

            valid_dates, valid_total_market_values, valid_circulating_market_values, valid_ratios, valid_listed_companies, valid_average_pe_ratios = process_data(rows)

            # 打印处理后的数据内容
            print(f"处理后 {exchange} 的有效日期数据: {valid_dates}")
            print(f"处理后 {exchange} 的有效总市值数据: {valid_total_market_values}")
            print(f"处理后 {exchange} 的有效流通市值数据: {valid_circulating_market_values}")
            print(f"处理后 {exchange} 的有效比例数据: {valid_ratios}")
            print(f"处理后 {exchange} 的有效上市公司数数据: {valid_listed_companies}")
            print(f"处理后 {exchange} 的有效平均市盈率数据: {valid_average_pe_ratios}")

            if not valid_dates:
                print(f"{exchange} 没有有效的数据用于绘图")
                continue

            all_valid_dates.extend(valid_dates)
            all_valid_total_market_values.extend(valid_total_market_values)
            all_valid_circulating_market_values.extend(valid_circulating_market_values)
            all_valid_ratios.extend(valid_ratios)
            all_valid_listed_companies.extend(valid_listed_companies)
            all_valid_average_pe_ratios.extend(valid_average_pe_ratios)

            if exchange == '上海交易所':
                valid_dates_shanghai = valid_dates
                valid_total_market_values_shanghai = valid_total_market_values
                valid_circulating_market_values_shanghai = valid_circulating_market_values
                valid_ratios_shanghai = valid_ratios
                valid_listed_companies_shanghai = valid_listed_companies
                valid_average_pe_ratios_shanghai = valid_average_pe_ratios
            elif exchange == '深圳交易所':
                valid_dates_shenzhen = valid_dates
                valid_total_market_values_shenzhen = valid_total_market_values
                valid_circulating_market_values_shenzhen = valid_circulating_market_values
                valid_ratios_shenzhen = valid_ratios
                valid_listed_companies_shenzhen = valid_listed_companies
                valid_average_pe_ratios_shenzhen = valid_average_pe_ratios

        # 绘制总市值和总流通市值图
        ax1 = plt.subplot(2, 2, 1)
        plot_subplot(ax1, valid_dates_shanghai, valid_total_market_values_shanghai, '上海交易所总市值', '总市值趋势', '日期', '市值（亿元）')
        plot_subplot(ax1, valid_dates_shanghai, valid_circulating_market_values_shanghai, '上海交易所流通市值', '总市值趋势', '日期', '市值（亿元）')
        plot_subplot(ax1, valid_dates_shenzhen, valid_total_market_values_shenzhen, '深圳交易所总市值', '总市值趋势', '日期', '市值（亿元）')
        plot_subplot(ax1, valid_dates_shenzhen, valid_circulating_market_values_shenzhen, '深圳交易所流通市值', '总市值趋势', '日期', '市值（亿元）')

        # 绘制流通市值占比图
        ax2 = plt.subplot(2, 2, 2)
        plot_subplot(ax2, valid_dates_shanghai, valid_ratios_shanghai, '上海交易所流通市值/总市值比例', '流通市值/总市值比例趋势', '日期', '比例')
        plot_subplot(ax2, valid_dates_shenzhen, valid_ratios_shenzhen, '深圳交易所流通市值/总市值比例', '流通市值/总市值比例趋势', '日期', '比例')

        # 绘制上市公司数图
        ax3 = plt.subplot(2, 2, 3)
        plot_subplot(ax3, valid_dates_shanghai, valid_listed_companies_shanghai, '上海交易所上市公司数', '上市公司数趋势', '日期', '公司数量')
        plot_subplot(ax3, valid_dates_shenzhen, valid_listed_companies_shenzhen, '深圳交易所上市公司数', '上市公司数趋势', '日期', '公司数量')

        # 绘制平均市盈率图
        ax4 = plt.subplot(2, 2, 4)
        plot_subplot(ax4, valid_dates_shanghai, valid_average_pe_ratios_shanghai, '上海交易所平均市盈率', '平均市盈率趋势', '日期', '市盈率')
        plot_subplot(ax4, valid_dates_shenzhen, valid_average_pe_ratios_shenzhen, '深圳交易所平均市盈率', '平均市盈率趋势', '日期', '市盈率')

        # 手动调整子图间距
        plt.subplots_adjust(left=0.1, right=0.95, top=0.95, bottom=0.1, hspace=0.3, wspace=0.2)

        # 之前使用的 plt.tight_layout() 可以注释掉，因为手动调整后可能不需要它了
        # plt.tight_layout()
        plt.show()

    except sqlite3.Error as e:
        print(f"数据库错误: {e}")
    finally:
        # 关闭数据库连接
        if conn:
            conn.close()

if __name__ == "__main__":
    plot_data_trend()